package demo;

import blockr.*;
import blockr.function.Output;
import blockr.function.RandomNum;
import blockr.model.RandomNumOBJ;
import usb.DecryptPrepare;
import usb.EncryptPrepare;
import util.CommonUtil;
import util.RSAUtil;

import java.io.File;
import java.io.IOException;
import java.io.UnsupportedEncodingException;

/**
 * @author spencercjh
 */
@SuppressWarnings("WeakerAccess")
public class RsaBlockR {
    /**
     * remoteDisk 存放相关文件的移动存储盘符
     * localDisk 存放相关文件的本地存储盘符
     * keyStore RSA公钥和私钥的根目录
     * publicKeystore 公钥文件路径
     * privateKeystore 私钥文件路径
     * inputFile 要加密的文件
     * encFile 原文Block-R算法加密后的密文
     * encFileCopy 原文Block-R算法加密后的密文拷贝
     * decFile 密文Block-R解密后的明文
     * brKeyFile(1/2) Block-R算法需要的秘钥原文
     * encBrKeyFile(1/2) Block-R算法秘钥RSA加密后的密文
     * decBrKeyFile(1/2) Block-R算法秘钥RSA解密后的明文
     * arraySize 目前的秘钥长度
     * charsetName 测试txt文件的编码
     */
    public static String remoteDisk = "D:\\test\\remote\\";
    public static String localDisk = "D:\\test\\local\\";

    public static String keyStore = remoteDisk + "RSA\\";
    public static String publicKeystore = remoteDisk + "RSA\\publicKey.keystore";
    public static String privateKeystore = remoteDisk + "RSA\\privateKey.keystore";

    public static String inputFile = localDisk + "small1.txt";
    public static String encFile = localDisk + "small1_enc.txt";
    public static String encFileCopy = remoteDisk + "RSA\\small1_enc.txt";
    public static String decFile = localDisk + "RSA\\small1_dec.txt";

    public static String brKeyFile1 = remoteDisk + "RSA\\brKeyFile1.txt";
    public static String brKeyFile2 = remoteDisk + "RSA\\brKeyFile2.txt";

    public static String encBrKeyFile1 = remoteDisk + "RSA\\encBrKeyFile1.txt";
    public static String encBrKeyFile2 = remoteDisk + "RSA\\encBrKeyFile2.txt";

    public static String decBrKeyFile1 = remoteDisk + "RSA\\decBrKeyFile1.txt";
    public static String decBrKeyFile2 = remoteDisk + "RSA\\decBrKeyFile2.txt";

    public static int arraySize = 128;
    public static String charsetName = "utf-8";

    public static void main(String[] args) {
//        encrypt(inputFile,false);
        decrypt("D:\\test\\local\\small1_enc.txt", false);
    }

    /***
     *
     * @param filePath 加密文件的路径
     * @param isUSBDevice 是否将密文密钥存储到U盘
     */
    private static void encrypt(String filePath, Boolean isUSBDevice) {
        if (!isUSBDevice) {
            /*秘钥存储路径*/
//            remoteDisk = remoteDisk + "test\\";
            // 得到要加密的文件路径，创建密文和秘钥的文件
            initDir(filePath);
//            init(remoteDisk);
            /*生成RSA的公钥和秘钥 路径分别为publicKeystore privateKeystore*/
            RSAUtil.generateKeyPair(keyStore);
            /*随机生成一对R1和R2作为Block-R加密的基础秘钥*/
            RandomNumOBJ firstBlockKey = new RandomNumOBJ(RandomNum.generate_random_num(arraySize), RandomNum.generate_random_num(arraySize));
            try {
                /*将R1和R2存入外存 路径分别为brKeyFile1 brKeyFile2*/
                CommonUtil.spitFile(brKeyFile1, firstBlockKey.getR1());
                CommonUtil.spitFile(brKeyFile2, firstBlockKey.getR2());
            } catch (IOException e) {
                e.printStackTrace();
                System.err.println("秘钥原文写失败");
            }

            /*对路径为inputfile的原文进行Block-R加密*/
            byte[][][] encBlockR = Encrypt.encrypt(inputFile, firstBlockKey);
            /*将Block-R加密过的密文存储到外存 路径为encFile和encFileCopy*/
            Output.output(encFile, encBlockR);
            Output.output(encFileCopy, encBlockR);
            try {
                /*读取公钥并RSA加密R1和R2*/
                String encbrKey1 = RSAUtil.fileEncrypt(publicKeystore, firstBlockKey.getR1());
                String encbrKey2 = RSAUtil.fileEncrypt(publicKeystore, firstBlockKey.getR2());
                /*将RSA加密后的R1和R2存储到外存 路径分别为encbrKeyFile1 encBrKeyFile2*/
                if (encbrKey1 == null || encbrKey2 == null) {
                    throw new AssertionError();
                }
                CommonUtil.spitFile(encBrKeyFile1, encbrKey1.getBytes(charsetName));
                CommonUtil.spitFile(encBrKeyFile2, encbrKey2.getBytes(charsetName));
            } catch (UnsupportedEncodingException e) {
                e.printStackTrace();
                System.err.println("RSA加密秘钥失败");
            } catch (IOException e) {
                e.printStackTrace();
                System.err.println("加密后秘钥写失败");
            }
        }
        /*加密前开始检测在10s内是否用户插入了秘钥U盘*/
        // U盘模式
        else if (isUSBDevice) {
            String newDevice = EncryptPrepare.init();
            if (!"error".equals(newDevice)) {
                /*修改系统变量*/
                remoteDisk = newDevice + "test\\";
                init(remoteDisk);
                /*生成RSA的公钥和秘钥 路径分别为publicKeystore privateKeystore*/
                RSAUtil.generateKeyPair(keyStore);
                /*随机生成一对R1和R2作为Block-R加密的基础秘钥*/
                RandomNumOBJ firstBlockKey = new RandomNumOBJ(RandomNum.generate_random_num(arraySize), RandomNum.generate_random_num(arraySize));
                try {
                    /*将R1和R2存入外存 路径分别为brKeyFile1 brKeyFile2*/
                    CommonUtil.spitFile(brKeyFile1, firstBlockKey.getR1());
                    CommonUtil.spitFile(brKeyFile2, firstBlockKey.getR2());
                } catch (IOException e) {
                    e.printStackTrace();
                    System.err.println("秘钥原文写失败");
                }

                /*对路径为inputfile的原文进行Block-R加密*/
                byte[][][] encBlockR = Encrypt.encrypt(inputFile, firstBlockKey);
                /*将Block-R加密过的密文存储到外存 路径为encFile和encFileCopy*/
                Output.output(encFile, encBlockR);
                Output.output(encFileCopy, encBlockR);
                try {
                    /*读取公钥并RSA加密R1和R2*/
                    String encbrKey1 = RSAUtil.fileEncrypt(publicKeystore, firstBlockKey.getR1());
                    String encbrKey2 = RSAUtil.fileEncrypt(publicKeystore, firstBlockKey.getR2());
                    /*将RSA加密后的R1和R2存储到外存 路径分别为encbrKeyFile1 encBrKeyFile2*/
                    if (encbrKey1 == null || encbrKey2 == null) {
                        throw new AssertionError();
                    }
                    CommonUtil.spitFile(encBrKeyFile1, encbrKey1.getBytes(charsetName));
                    CommonUtil.spitFile(encBrKeyFile2, encbrKey2.getBytes(charsetName));
                } catch (UnsupportedEncodingException e) {
                    e.printStackTrace();
                    System.err.println("RSA加密秘钥失败");
                } catch (IOException e) {
                    e.printStackTrace();
                    System.err.println("加密后秘钥写失败");
                }
            } else {
                System.err.println("没有找到秘钥U盘");
            }
        }

    }

    /***
     *
     * @param filePath 密文地址
     * @param isUSBDevice 是否U盘设备，是则自动在U盘中扫描密钥文件
     */

    private static void decrypt(String filePath, boolean isUSBDevice) {
        if (!isUSBDevice) {
            RandomNumOBJ firstBlockKey = new RandomNumOBJ();
            initDir(filePath);
            try {
                /*读取RSA加密过的R1和R2 路径分别为encbrKeyFile1 encBrKeyFile2*/
                String decbrKey1 = new String(CommonUtil.suckFile(encBrKeyFile1), charsetName);
                String decbrKey2 = new String(CommonUtil.suckFile(encBrKeyFile2), charsetName);
                /*读取私钥解密R1和R2*/
                firstBlockKey = new RandomNumOBJ();
                firstBlockKey.setR1(RSAUtil.fileDecrypt(privateKeystore, decbrKey1));
                firstBlockKey.setR2(RSAUtil.fileDecrypt(privateKeystore, decbrKey2));
            } catch (IOException e) {
                e.printStackTrace();
                System.err.println("加密后秘钥读失败");
            }
            try {
                /*将解密后的R1和R2存储到外存 路径为decbrKeyFile1 decBrKeyFile2*/
                CommonUtil.spitFile(decBrKeyFile1, firstBlockKey.getR1());
                CommonUtil.spitFile(decBrKeyFile2, firstBlockKey.getR2());
            } catch (IOException e) {
                e.printStackTrace();
                System.err.println("秘钥解密后原文写失败");
            }
            /*将U盘中的密文与本地存储的密文进行比较,如果不同进行覆盖*/
            try {
                if (!CommonUtil.issame(encFile, encFileCopy)) {
                    byte[] correctContent = CommonUtil.suckFile(encFileCopy);
                    CommonUtil.spitFile(encFile, correctContent);
                    System.out.println("正确密文覆盖成功！");
                }
            } catch (IOException e) {
                e.printStackTrace();
                System.err.println("正确密文覆盖失败");
            }
            /*利用解密出来的R1和R2解密原文件 路径为encfile*/
            byte[][][] decBlockR = Decrypt.decrypt(encFile, arraySize, firstBlockKey);
            /*将解密后的明文存储到外存 路径为decfile*/
            Output.outputContent(decFile, decBlockR, charsetName);
            System.out.println("success over!!!");
        }
        // U盘模式
        else if (isUSBDevice) {
            String newDevice = DecryptPrepare.init();
            if (!"error".equals(newDevice)) {
                /*修改系统变量*/
                remoteDisk = newDevice + "test\\";
            }
            RandomNumOBJ firstBlockKey = new RandomNumOBJ();
            try {
                /*读取RSA加密过的R1和R2 路径分别为encbrKeyFile1 encBrKeyFile2*/
                String decbrKey1 = new String(CommonUtil.suckFile(encBrKeyFile1), charsetName);
                String decbrKey2 = new String(CommonUtil.suckFile(encBrKeyFile2), charsetName);
                /*读取私钥解密R1和R2*/
                firstBlockKey = new RandomNumOBJ();
                firstBlockKey.setR1(RSAUtil.fileDecrypt(privateKeystore, decbrKey1));
                firstBlockKey.setR2(RSAUtil.fileDecrypt(privateKeystore, decbrKey2));
            } catch (IOException e) {
                e.printStackTrace();
                System.err.println("加密后秘钥读失败");
            }
            try {
                /*将解密后的R1和R2存储到外存 路径为decbrKeyFile1 decBrKeyFile2*/
                CommonUtil.spitFile(decBrKeyFile1, firstBlockKey.getR1());
                CommonUtil.spitFile(decBrKeyFile2, firstBlockKey.getR2());
            } catch (IOException e) {
                e.printStackTrace();
                System.err.println("秘钥解密后原文写失败");
            }
            /*将U盘中的密文与本地存储的密文进行比较,如果不同进行覆盖*/
            try {
                if (!CommonUtil.issame(encFile, encFileCopy)) {
                    byte[] correctContent = CommonUtil.suckFile(encFileCopy);
                    CommonUtil.spitFile(encFile, correctContent);
                    System.out.println("正确密文覆盖成功！");
                }
            } catch (IOException e) {
                e.printStackTrace();
                System.err.println("正确密文覆盖失败");
            }
            /*利用解密出来的R1和R2解密原文件 路径为encfile*/
            byte[][][] decBlockR = Decrypt.decrypt(encFile, arraySize, firstBlockKey);
            /*将解密后的明文存储到外存 路径为decfile*/
            Output.outputContent(decFile, decBlockR, charsetName);
            System.out.println("success over!!!");
        } else {
            System.err.println("没有找到秘钥U盘");
        }
    }

    private static void init(String remoteDisk) {
        keyStore = remoteDisk + "RSA\\";
        publicKeystore = remoteDisk + "RSA\\publicKey.keystore";
        privateKeystore = remoteDisk + "RSA\\privateKey.keystore";
        encFileCopy = remoteDisk + "RSA\\small1_enc.txt";
        brKeyFile1 = remoteDisk + "RSA\\brKeyFile1.txt";
        brKeyFile2 = remoteDisk + "RSA\\brKeyFile2.txt";
        encBrKeyFile1 = remoteDisk + "RSA\\encBrKeyFile1.txt";
        encBrKeyFile2 = remoteDisk + "RSA\\encBrKeyFile2.txt";
        decBrKeyFile1 = remoteDisk + "RSA\\decBrKeyFile1.txt";
        decBrKeyFile2 = remoteDisk + "RSA\\decBrKeyFile2.txt";
    }

    // 初始化密文和秘钥存储位置
    private static void initDir(String filePath) {
        File file = new File(filePath);
        String dir = file.getParent();// 父路径
        String encFileDir = dir + "\\encFile";// 密文文件夹路径
        keyStore = dir + "\\keyStore"; // 密钥文件夹路径
        String decFileDir = dir + "\\decFile"; // 解密后文件夹
        // 创建文件夹
        File decFileDirFile = new File(decFileDir);
        if (!decFileDirFile.exists()) {
            decFileDirFile.mkdirs();
        }
        File encfileDirFile = new File(encFileDir);
        if (!encfileDirFile.exists()) {
            encfileDirFile.mkdirs();
        }
        File keyStoreFile = new File(keyStore);
        if (!keyStoreFile.exists()) {
            keyStoreFile.mkdirs();
        }

        publicKeystore = keyStore + "\\publicKey.keystore"; // 公钥文件路径
        privateKeystore = keyStore + "\\privateKey.keystore"; // 私钥文件路径
        encFile = encFileDir + "\\small1_enc.txt"; // 密文文件路径
        encFileCopy = encFileDir + "\\small1_enc.txt"; // 密文文件路径
        decFile = decFileDir + "\\small1_dec.txt"; // 解密后文件路径

        // br秘钥文件路径
        brKeyFile1 = keyStore + "\\brKeyFile1.txt";
        brKeyFile2 = keyStore + "\\brKeyFile2.txt";
        encBrKeyFile1 = keyStore + "\\encBrKeyFile1.txt";
        encBrKeyFile2 = keyStore + "\\encBrKeyFile2.txt";

        decBrKeyFile1 = keyStore + "\\decBrKeyFile1.txt";
        decBrKeyFile2 = keyStore + "\\decBrKeyFile2.txt";
    }

}
